package io.github.fakeport;


import java.io.Closeable;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;

public class TcpForward implements Runnable ,Closeable{
    private static int BUF_SIZE = 4096;
    private Selector selector;

    private ServerSocketChannel serverSocket=null;

    // write queue
    private HashMap<SelectionKey,ByteBuffer> writeQueue = new HashMap<>();
    // byte buffer pool
    private LinkedList<ByteBuffer> bufferPool = new LinkedList<>();

    public TcpForward(int local_port) throws IOException {
        this.listen(local_port);
    }

    /**
     * listen local port and create selector
     * @param local_port local port
     * @return boolean success
     * @throws IOException System Error
     */
    private void listen(int local_port) throws IOException {
        selector = Selector.open();
        // Open channel
        serverSocket = ServerSocketChannel.open();
        // bind local port
        try {
            //serverSocket.bind(new InetSocketAddress(local_port),10);
            serverSocket.configureBlocking(false);
            serverSocket.socket().bind(new InetSocketAddress(local_port));
            serverSocket.socket().setReuseAddress(true);
            serverSocket.register(selector, SelectionKey.OP_ACCEPT);
            //return true;
        } catch (IOException e) {
            serverSocket.close();
            selector.close();
            throw e;
        }
        //return false;
    }

    @Override
    public void run() {
        int err_select = 0;
        while (true){
            // wait client connect
            try {
                selector.select();
            } catch (IOException e) {
                e.printStackTrace();
                if (err_select++ > 0)
                    break;
                else
                    continue;
            }
            err_select = 0;
            // process event
            Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
            while (iterator.hasNext()){ // foreach key
                SelectionKey key = iterator.next();
                iterator.remove();
                try {
                    // event channel is server socket
                    if (key.isAcceptable()) {
                        onAcceptable(key);
                    } else if (key.isReadable()) {
                        onReadable(key);
                    } else if (key.isConnectable()) {
                        onConnectable(key);
                    } else if (key.isWritable()) {
                        onWriteable(key);
                    }
                }catch (CancelledKeyException e){
                    e.printStackTrace();
                    //Log.e(TAG, "run: event", e);
                    //key.
                }
            }


        }
    }

    /**
     * process write queue and register read event
     * @param key event key
     */
    private void onWriteable(SelectionKey key) {
        // socket (can write)
        SocketChannel sock = (SocketChannel) key.channel();

        if (writeQueue.containsKey(key)){
            ByteBuffer buffer = writeQueue.get(key);
            int bywrite = 0;
            try {
                // write data to socket(from buffer queue)
                while (buffer.hasRemaining()){
                    bywrite += sock.write(buffer);
                }
                writeQueue.remove(key);
                // set it unused
                freeBuffer(buffer);
            } catch (IOException e) {
                e.printStackTrace();
                if (bywrite <= 0)
                    // 这里为什么要监听read呢。。。
                    // 很简单，如果我在这里close掉fd
                    // 那就需要把下面的代码上移，并且会加大这个函数的复杂程度
                    // 让他触发READ事件，然后再关闭比较好
                    key.interestOps(SelectionKey.OP_READ);
                return ;
            }

        }

        SelectionKey other_key = (SelectionKey) key.attachment();

        // change event (write finished)
        key.interestOps(key.interestOps() &(~SelectionKey.OP_WRITE)); // unregister write event
        other_key.interestOps(key.interestOps() | SelectionKey.OP_READ); // register read event(wait data)
    }

    /**
     * finish connect
     * listen read event
     * @param key event key
     */
    private void onConnectable(SelectionKey key) {
        SocketChannel sock = (SocketChannel) key.channel();
        SelectionKey other_key = (SelectionKey) key.attachment();
        try {
            sock.finishConnect();
            key.interestOps(SelectionKey.OP_WRITE);
            other_key.interestOps(SelectionKey.OP_WRITE);
            System.out.println("Connected");
        } catch (IOException e) {
            e.printStackTrace();
            closeSock(key);
            /*sock.close();
            other_key.channel().close();*/
        }
    }

    /**
     * forward data and or write queue
     * @param key event key
     */
    private void onReadable(SelectionKey key) {
        SocketChannel sock = (SocketChannel) key.channel();
        SelectionKey other_key = (SelectionKey) key.attachment();

        ByteBuffer buffer = null;
        boolean dontFreeBuffer = false;
        try {
            if (other_key == null) {
                // Connect by fake port
                ByteBuffer fb = ByteBuffer.allocate(4);// int
                sock.read(fb);
                fb.flip();
                if (fb.limit() != fb.capacity())
                    throw new Exception("Can read fake port !");

                int port = fb.getInt();
                if (port <= 0 || port > 65535 || port == serverSocket.socket().getLocalPort())
                    throw new Exception("Fake port " + port + " is not supported ! ");
                System.out.println("Connect fake port: " + port);
                SocketChannel other = SocketChannel.open();
                other.configureBlocking(false);
                other.connect(new InetSocketAddress("127.0.0.1", port));
                key.interestOps(0);
                other_key = other.register(selector, SelectionKey.OP_CONNECT);
                other_key.attach(key);
                key.attach(other_key);
            } else {
                // forward
                // alloc buffer
                SocketChannel other = (SocketChannel) other_key.channel();
                buffer = allocBuffer();
                //try {
                while (true) {
                    buffer.clear(); // clear buffer
                    // try to read data
                    int byread = sock.read(buffer);
                    // if socket is closed
                    if (byread < 0)
                        throw new IOException("Can't read from socket, because it closed ! " + sock.socket().getRemoteSocketAddress().toString());
                    // if read finished
                    if (byread == 0)
                        break;

                    buffer.flip();
                    int bywrite = 0,tmp;
                    try {
                        while (buffer.hasRemaining()) {
                            tmp = other.write(buffer);
                            if (tmp < 1)
                                throw new IOException("Write zero bytes");
                            bywrite += tmp;
                        }
                    } catch (IOException e) {
                        // write error
                        if (bywrite <= 0) {
                            throw e;
                        }
                        // add write queue
                        writeQueue.put(other_key, buffer);
                        dontFreeBuffer = true;
                        // change event
                        key.interestOps(key.interestOps() & (~SelectionKey.OP_READ));// unregister read event
                        other_key.interestOps(other_key.interestOps() | SelectionKey.OP_WRITE); // register remote write event
                        break; // ignore data
                    }


                }
            }
        }catch (Exception e){
            e.printStackTrace();
            closeSock(key);
        }
        if (buffer!=null && !dontFreeBuffer)
            freeBuffer(buffer);
    }

    /**
     * accept client socket and connect to remote server
     * create selection key of all socket
     * listen connect event
     * @param key event key
     */
    private void onAcceptable(SelectionKey key){
        ServerSocketChannel serverSocket = (ServerSocketChannel) key.channel();
        SocketChannel channel = null;
        try {
            channel = serverSocket.accept();
            System.out.println("Accept new client: " + channel.toString());
            channel.configureBlocking(false);
            // register none event
            channel.register(key.selector(), SelectionKey.OP_READ);
        } catch (IOException e) {
            if (channel != null)
                try {
                    channel.close();
                } catch (IOException e1) {
                    //e1.printStackTrace();
                }
            e.printStackTrace();
        }
    }

    /**
     * close two socket (^_^)
     * @param key key to close
     */
    private void closeSock(SelectionKey key, boolean closeBorther){
        System.out.println("closeSock: close a socket");
        cleanWriteQueueBySock(key);
        try {
            key.channel().close();
            if (closeBorther && key.attachment() != null) {
                closeSock((SelectionKey) key.attachment(),false);
            }
        }catch (IOException ignored){

        }
    }

    private void closeSock(SelectionKey key){
        closeSock(key,true);
    }

    private void cleanWriteQueueBySock(SelectionKey key){
        if (writeQueue.containsKey(key))
            freeBuffer(writeQueue.remove(key));
    }

    /**
     * alloc buffer from pool
     * @return ByteBuffer
     */
    private ByteBuffer allocBuffer(){
        if (bufferPool.isEmpty()) {
            System.out.println("allocBuffer: alloc new buffer");
            return ByteBuffer.allocate(BUF_SIZE);
        }
        //return bufferPool.poll();
        return bufferPool.removeFirst();
    }

    /**
     * free buffer
     * @param buffer to free
     */
    private void freeBuffer(ByteBuffer buffer){
        buffer.clear();
        //bufferPool.offer(buffer);
        bufferPool.addFirst(buffer);
    }



    @Override
    public void close() /*throws IOException*/ {
        // stop select
        if (selector != null) {
            selector.wakeup();
            try {
                selector.close();
            } catch (IOException ignored) {
            }
            selector = null;
        }
        if (serverSocket != null){
            try {
                serverSocket.close();
            }catch (IOException ignored){

            }
        }
    }

}
